import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.InputStreamReader;
import java.text.DecimalFormat;
import java.text.NumberFormat;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;

/**
 * GetFastaStas computes various statistics about the sequences in fasta files.
 *
 * The parameter min is used to set a lower bound for contig length. Contigs
 * with lengths lower than the minimum are ignored.
 *
 * The following stats are computed:
 * Total units: The number of sequences/contigs in the input files.
 * Reference: The size of the reference genome. If the argument
 *   genomesize is provided then this is the value of that argument.
 *   If genomesize is not provided this is the sum of the contig lengths.
 * BasesInFasta: This is the sum of the lengths of the contigs. Only
 *   contigs with a length larger than the minimum are included.
 * Min: The length of the shortest contig.
 * Max: The length of the longest contig.
 * N10, N25, N75: These are the lengths of the contigs such that the
 *   sum of the contigs longer than this length is = X% of the total genome.
 *   The percentage is computed with respect to the value of Reference.
 *   The count is the number of sequences with length longer than this length.
 * E-size: This is a weighted average of the contig lengths. The weight
 *   is contig_length/reference.
 *
 * TODO(jeremy@lewi.us): getCorrectness stats is always using option -o
 * so we should get rid of support for the other formats.
 */
public class GetFastaStats {
   private static final int MIN_GAP_SIZE = 20;
   private static int MIN_LENGTH = 2000;
   private static final int CONTIG_AT_INITIAL_STEP = 1000000;

   private static final String[] suffix = {
     "final", "split", "fasta", "scafSeq", "fa"};
   private static final NumberFormat nf =
       new DecimalFormat("############.##");

   private class ContigAt {
	   ContigAt(long currentBases) {
		   this.count = this.len = 0;
		   this.totalBP = 0;
		   this.goal = currentBases;
	   }

	   public int count;
	   public long totalBP;
	   public int len;
	   public long goal;
   }

   boolean baylorFormat = false;
   boolean oldStyle = false;
   boolean storeCtgs = false;
   HashMap<String, Integer> fastaLens = new HashMap<String, Integer>();
   HashMap<String, String> ctgs = new HashMap<String, String>();
   int totalCount = 0;
   private long genomeSize = 0;

   public GetFastaStats(boolean useBaylor, boolean old, long g) {
	   baylorFormat = useBaylor;
	   oldStyle = old;
           genomeSize = g;
   }

   // (jeremy@lewi.us): Looks like this function is loading the lengths
   // of the sequences into the hasmap fastaLens. The input can either
   // be a fasta file or a file with "lens" in its name in which case
   // it contains the lengths of the various contigs.
   public void processFile(String inputFile) throws Exception {
     if (inputFile.contains("lens")) {
       BufferedReader bf = Utils.getFile(inputFile, "lens");
       String line = null;
       while ((line = bf.readLine()) != null) {
         String[] splitLine = line.trim().split("\\s+");
         fastaLens.put(splitLine[0], Integer.parseInt(splitLine[1]));
       }
       bf.close();
     } else {
       BufferedReader bf = Utils.getFile(inputFile, suffix);

       String line = null;
       StringBuffer fastaSeq = new StringBuffer();
       String header = "";

       while ((line = bf.readLine()) != null) {
         if (line.startsWith(">")) {
           String fastaString = fastaSeq.toString().replaceAll("-", "");
           fastaString = fastaString.toString().replaceAll("\\.", "");

           //String[] split = fastaString.trim().split("N+");
           //String[] split = { fastaString.replaceAll("N", "").trim() };
           String[] split = { fastaString.trim() };
           //System.err.println("SPLIT one ctg of length " + fastaString.length() + " INTO " + split.length);

           for (int i = 0; i < split.length; i++) {
             if (split[i].length() != 0) {
               fastaLens.put(header + "_" + i,split[i].length());
               if (storeCtgs)
                 ctgs.put(header + "_" + i, split[i].toString());
             }
           }
           //header = line.replace(">contig_", "").trim();
           header = line.trim().split("\\s+")[0].replace(">", "").replaceAll("ctg", "").trim();
           header = line.trim().split("\\s+")[0].replace(">", "").replaceAll("scf", "").trim();
           header = line.trim().split("\\s+")[0].replace(">", "").replaceAll("deg", "").trim();
           fastaSeq = new StringBuffer();
         }
         else {
           fastaSeq.append(line);
         }
       }

       String fastaString = fastaSeq.toString().replaceAll("-", "");
       fastaString = fastaString.toString().replaceAll("\\.", "");

       //String[] split = fastaString.trim().split("N+");
       //System.err.println("SPLIT one ctg of length " + fastaString.length() + " INTO " + split.length);
       //String[] split = { fastaString.replaceAll("N", "").trim() };
       String[] split = { fastaString.trim() };
       for (int i = 0; i < split.length; i++) {
         if (fastaSeq.length() != 0) {
           fastaLens.put(header + "_" + i,split[i].length());

           if (storeCtgs)
             ctgs.put(header + "_" + i, split[i].toString());
         }
       }
       bf.close();
     }
   }

   // TODO(jeremy@lewi.us): countReads is never called because the invocation
   // in main is commented out. So we should get read of this function.
   public void countReads(String posmapFile, boolean isNB) throws Exception {
	      BufferedReader bf = new BufferedReader(new InputStreamReader(
	              new FileInputStream(posmapFile)));

       String line = null;
       String currID = "";
       int readCount = 0;
       totalCount = 0;
       int contigIDIndex = (isNB ? 2 : 1);

       while ((line = bf.readLine()) != null) {
    	   String[] splitLine = line.trim().split("\\s+");

    	   if (splitLine.length <= contigIDIndex) {
    		   continue;
    	   }
    	   String id = splitLine[contigIDIndex].trim();

    	   if (!id.equalsIgnoreCase(currID)) {
    		   totalCount += readCount;
    		   currID = id;
    		   readCount = 0;
    	   }
    	   if (fastaLens.get(id) != null) {
    		   readCount++;
    	   }
       }
       if (readCount != 0) {
    	   totalCount += readCount;
       }
   }

   // TODO(jeremy@lewi.us): convertToScaffolds is never called because the
   // invocation in main is commented out. So we should get read of this
   // function.
   public void convertToScaffolds(String scaffFile) throws Exception {
      BufferedReader bf = new BufferedReader(new InputStreamReader(
            new FileInputStream(scaffFile)));

      String line = null;
      String currID = "";
      StringBuffer fastaSeq = new StringBuffer();
      int negGaps = 0;
      String id = "";

      while ((line = bf.readLine()) != null) {
         String[] splitLine = line.trim().split("\\s+");

         try {
            if (splitLine[0].equalsIgnoreCase("supercontig")) {
               id = splitLine[1];
            }
            if (splitLine[0].equalsIgnoreCase("contig")) {
               fastaSeq.append(ctgs.get(splitLine[1]));
            }
            if (splitLine[0].equalsIgnoreCase("gap")) {
               int gapSize = Integer.parseInt(splitLine[1]);

               if (gapSize < 0) {
                  negGaps += gapSize;
                  fastaSeq.delete(fastaSeq.length()+gapSize, fastaSeq.length());
                  //gapSize = MIN_GAP_SIZE;
               }
               for (int i = 0; i < gapSize; i++) {
                  fastaSeq.append("N");
               }
            }
         } catch (Exception e) {
            continue;
         }

         if (!currID.equalsIgnoreCase(id)) {
            if (fastaSeq.length() != 0) {
               System.err.println(">scaffold_" + currID);
               System.err.println(Utils.convertToFasta(fastaSeq.toString().trim()));
               fastaSeq = new StringBuffer();
            }
            currID = id;
         }
      }

      if (fastaSeq.length() != 0) {
         System.err.println(">scaffold_" + currID);
         System.err.println(Utils.convertToFasta(fastaSeq.toString().trim()));
      }

      System.out.println("NEGATIVE GAPS IS " + negGaps);
   }

   /**
    * Compute the statistics and convert them to a string.
    * @param outputHeader
    * @param title
    * @return
    */
   public String toString(boolean outputHeader, String title) {
     StringBuffer st = new StringBuffer();
     int max = Integer.MIN_VALUE;
     int min = Integer.MAX_VALUE;
     long total = 0;
     int count = 0;

     int n10 = 0;
     int n25 = 0;
     int n50 = 0;
     int n75 = 0;
     int n95 = 0;

     int n10count = 0;
     int n25count = 0;
     int n50count = 0;
     int n75count = 0;
     int n95count = 0;

     int totalOverLength = 0;
     long totalBPOverLength = 0;

     double fSize = 0;
     for (String s : fastaLens.keySet()) {
       int len = fastaLens.get(s);

       if (len <= MIN_LENGTH) {
         continue;
       }
       if (oldStyle == true && len <= MIN_LENGTH) {
         continue;
       }

       if (len > max) { max = len; }
       if (len < min) { min = len; }
       if (genomeSize == 0) {
         total += len;
       } else {
         total = genomeSize;
       }
       count++;

       // compute the F-size
       fSize += Math.pow(len, 2);
       if (len > MIN_LENGTH) {
         totalOverLength++;
         totalBPOverLength += len;
       }
     }
     fSize /= genomeSize;

     // get the goal contig at X bases (1MBp, 2MBp)
     ArrayList<ContigAt> contigAtArray = new ArrayList<ContigAt>();
     if (baylorFormat == true) {
       contigAtArray.add(new ContigAt( 1 * CONTIG_AT_INITIAL_STEP));
       contigAtArray.add(new ContigAt( 2 * CONTIG_AT_INITIAL_STEP));
       contigAtArray.add(new ContigAt( 5 * CONTIG_AT_INITIAL_STEP));
       contigAtArray.add(new ContigAt(10 * CONTIG_AT_INITIAL_STEP));
     }
     else {
       long step = CONTIG_AT_INITIAL_STEP;
       long currentBases = 0;
       while (currentBases <= total) {
         if ((currentBases / step) >= 10) {
           step *= 10;
         }
         currentBases += step;
         contigAtArray.add(new ContigAt(currentBases));
       }
     }
     ContigAt[] contigAtVals = contigAtArray.toArray(new ContigAt[0]);

     Integer[] vals = fastaLens.values().toArray(new Integer[0]);
     Arrays.sort(vals);

     long sum = 0;
     double median = 0;
     int medianCount = 0;
     int numberContigsSeen = 1;
     int currentValPoint = 0;
     for (int i = vals.length - 1; i >= 0; i--) {
       if ((count / 2) == i) {
         median += vals[i];
         medianCount++;
       }
       else if (count % 2 == 0 && (((count / 2) + 1) == i)) {
         median += vals[i];
         medianCount++;
       }

       sum += vals[i];

       // calculate the bases at
       while (currentValPoint < contigAtVals.length && sum >= contigAtVals[currentValPoint].goal && contigAtVals[currentValPoint].count == 0) {
         System.err.println("Calculating point at " + currentValPoint + " and the sum is " + sum + " and i is" + i + " and lens is " + vals.length + " and length is " + vals[i]);
         contigAtVals[currentValPoint].count = numberContigsSeen;
         contigAtVals[currentValPoint].len = vals[i];
         contigAtVals[currentValPoint].totalBP = sum;
         currentValPoint++;
       }
       // calculate the NXs
       if (sum / (double)total >= 0.1 && n10count == 0) {
         n10 = vals[i];
         n10count = vals.length - i;

       }
       if (sum / (double)total >= 0.25 && n25count == 0) {
         n25 = vals[i];
         n25count = vals.length - i;

       }
       if (sum / (double)total >= 0.5 && n50count == 0) {
         n50 = vals[i];
         n50count = vals.length - i;

       }
       if (sum / (double)total >= 0.75 && n75count == 0) {
         n75 = vals[i];
         n75count = vals.length - i;
       }
       if (sum / (double)total >= 0.95 && n95count == 0) {
         n95 = vals[i];
         n95count = vals.length - i;

       }

       numberContigsSeen++;
     }
     if (medianCount != 1 && medianCount != 2) {
       System.err.println("ERROR INVALID MEDIAN COUNT " + medianCount);
       System.exit(1);
     }

     // TODO(jeremy@lewi.us): getCorrectnessStats.sh is always using
     // the oldStyle so we should get rid of support for the other fromats.
     if (oldStyle == true) {
       st.append("Total units: " + count + "\n");
       st.append("Reference: " + total + "\n");
       st.append("BasesInFasta: " + totalBPOverLength + "\n");
       st.append("Min: " + min + "\n");
       st.append("Max: " + max + "\n");
       st.append("N10: " + n10 + " COUNT: " + n10count + "\n");
       st.append("N25: " + n25 + " COUNT: " + n25count + "\n");
       st.append("N50: " + n50 + " COUNT: " + n50count + "\n");
       st.append("N75: " + n75 + " COUNT: " + n75count + "\n");
       st.append("E-size:" + nf.format(fSize) + "\n");
     } else {
       if (outputHeader) {
         st.append("Assembly");
         st.append(",Unit Number");
         st.append(",Unit Total BP");
         st.append(",Number Units > " + MIN_LENGTH);
         st.append(",Total BP in Units > " + MIN_LENGTH);
         st.append(",Min");
         st.append(",Max");
         st.append(",Average");
         st.append(",Median");

         for (int i = 0; i < contigAtVals.length; i++) {
           if (contigAtVals[i].count != 0) {
             st.append(",Unit At " + nf.format(contigAtVals[i].goal) + " Unit Count," + /*"Total Length," + */ " Actual Unit Length" );
           }
         }
         st.append("\n");
       }

       st.append(title);
       st.append("," + nf.format(count));
       st.append("," + nf.format(total));
       st.append("," + nf.format(totalOverLength));
       st.append("," + nf.format(totalBPOverLength));
       st.append("," + nf.format(min));
       st.append("," + nf.format(max));
       st.append("," + nf.format((double)total / count));
       st.append("," + nf.format(median / medianCount));

       for (int i = 0; i < contigAtVals.length; i++) {
         if (contigAtVals[i].count != 0) {
           st.append("," + contigAtVals[i].count + "," + /*contigAtVals[i].totalBP + "," +*/ contigAtVals[i].len);
         }
       }
     }

     return st.toString();
   }

   public static void printUsage() {
      System.err.println("This program computes total bases in fasta sequences, N50, min, and max.");

   }
   public static void main(String[] args) throws Exception {
      if (args.length < 1) { printUsage(); System.exit(1);}

      /*
      if (args.length > 1) {
         f.storeCtgs = true;
      }
      */

      boolean useBaylorFormat = false;
      boolean oldStyle = false;
      long genomeSize = 0;
      int initialVal = 0;
      while (args[initialVal].startsWith("-")) {
    	  if (args[initialVal].trim().equalsIgnoreCase("-b")) {
    		  useBaylorFormat = true;
    	  }
          else if (args[initialVal].trim().equalsIgnoreCase("-min")) {
                  GetFastaStats.MIN_LENGTH = Integer.parseInt(args[++initialVal]);
          }
    	  else if (args[initialVal].trim().equalsIgnoreCase("-o")) {
    		  oldStyle = true;
    	  }
          else if (args[initialVal].trim().equalsIgnoreCase("-genomeSize")) {
                  initialVal++;
                  genomeSize = Long.parseLong(args[initialVal]);
System.err.println("Found genome size at position " + initialVal + " with value " + args[initialVal] + " aka " + genomeSize);
          }
    	  else {
    		  System.err.println("Unknown parameter "+ args[initialVal] + " specified, please specify -b for Baylor-style output.");
    		  System.exit(1);
    	  }
    	  initialVal++;
      }

      for (int i = initialVal; i < args.length; i++) {
    	  String assemblyTitle = args[i].trim().split("/")[0];
          GetFastaStats f = new GetFastaStats(useBaylorFormat, oldStyle, genomeSize);
          String[] splitLine = args[i].trim().split(",");

          for (int j = 0; j < splitLine.length; j++) {
        	  f.processFile(splitLine[j]);
    	  }
          System.out.println(f.toString(i == initialVal, assemblyTitle));
      }
      /*
      if (args.length > 1) {
         //f.convertToScaffolds(args[1]);
    	  f.countReads(args[1], Boolean.parseBoolean(args[2]));
      }
      */
   }
}
